import json
from enum import IntEnum
from typing import Any, Callable, List, Optional

from lagent.prompts.parsers import StrParser
from lagent.utils import create_object, load_class_from_string


def default_plugin_validate(plugin: str):
    """
    默认的插件验证函数
    
    验证插件字符串是否为有效的JSON格式
    
    参数:
        plugin (str): 待验证的插件字符串
        
    返回:
        dict: 解析后的JSON对象
        
    异常:
        JSONDecodeError: 当插件字符串不是有效的JSON格式时抛出
    """
    plugin = plugin.strip()
    if not (plugin.startswith('{') and plugin.endswith("}")):
        raise json.decoder.JSONDecodeError
    return json.loads(plugin)


class ToolStatusCode(IntEnum):
    """
    工具状态码枚举类
    
    NO_TOOL: 没有使用工具
    VALID_TOOL: 有效的工具调用
    PARSING_ERROR: 解析错误
    """
    NO_TOOL = 0
    VALID_TOOL = 1
    PARSING_ERROR = -1


class ToolParser(StrParser):
    """
    工具解析器基类
    
    用于解析和格式化工具相关的文本内容
    
    参数:
        tool_type (str): 工具类型
        template (str): 模板字符串
        begin (str): 工具调用开始标记
        end (str): 工具调用结束标记
        validate (Callable): 验证函数
    """

    def __init__(
        self,
        tool_type: str,
        template: str = '',
        begin: str = '<tool>\n',
        end: str = '</tool>\n',
        validate: Callable[[str], Any] = None,
        **kwargs
    ):
        super().__init__(template, begin=begin, end=end, **kwargs)
        self.template = template
        self.tool_type = tool_type
        self.validate = load_class_from_string(validate) if isinstance(validate, str) else validate

    def parse_response(self, data: str) -> dict:
        """
        解析响应文本
        
        参数:
            data (str): 待解析的响应文本
            
        返回:
            dict: 包含工具类型、思考过程、动作和状态的字典
        """
        # 如果没有工具调用标记，返回普通文本响应
        if self.format_field['begin'] not in data:
            return dict(tool_type=None, thought=data, action=None, status=ToolStatusCode.NO_TOOL)
        
        # 分割思考过程和动作
        thought, action, *_ = data.split(self.format_field["begin"])
        action = action.split(self.format_field['end'])[0]
        status = ToolStatusCode.VALID_TOOL
        
        # 如果有验证函数，验证动作的有效性
        if self.validate:
            try:
                action = self.validate(action)
            except Exception:
                status = ToolStatusCode.PARSING_ERROR
        return dict(tool_type=self.tool_type, thought=thought, action=action, status=status)

    def format_response(self, parsed: dict) -> str:
        """
        格式化响应
        
        参数:
            parsed (dict): 解析后的响应字典
            
        返回:
            str: 格式化后的响应文本
        """
        if parsed['action'] is None:
            return parsed['thought']
        assert parsed['tool_type'] == self.tool_type
        
        # 将动作转换为字符串格式
        if isinstance(parsed['action'], dict):
            action = json.dumps(parsed['action'], ensure_ascii=False)
        else:
            action = str(parsed['action'])
        return parsed['thought'] + self.format_field['begin'] + action + self.format_field['end']


class InterpreterParser(ToolParser):
    """
    解释器解析器
    
    用于解析和格式化代码解释器相关的内容
    """
    def __init__(
        self,
        tool_type: str = 'interpreter',
        template: str = '',
        begin: str = '<|action_start|><|interpreter|>\n',
        end: str = '<|action_end|>\n',
        validate: Callable[[str], Any] = None,
        **kwargs
    ):
        super().__init__(tool_type, template, begin, end, validate, **kwargs)


class PluginParser(ToolParser):
    """
    插件解析器
    
    用于解析和格式化插件相关的内容，默认使用JSON格式验证
    """
    def __init__(
        self,
        tool_type: str = 'plugin',
        template: str = '',
        begin: str = '<|action_start|><|plugin|>\n',
        end: str = '<|action_end|>\n',
        validate: Callable[[str], Any] = default_plugin_validate,
        **kwargs
    ):
        super().__init__(tool_type, template, begin, end, validate, **kwargs)


class MixedToolParser(StrParser):
    """
    混合工具解析器
    
    支持多种工具类型的解析和格式化
    
    参数:
        tool_type (Optional[str]): 默认工具类型
        template (str): 模板字符串
        parsers (List[ToolParser]): 工具解析器列表
    """
    def __init__(self, tool_type: Optional[str] = None, template='', parsers: List[ToolParser] = None, **format_field):
        self.parsers = {}
        self.tool_type = tool_type
        # 初始化所有解析器
        for parser in parsers or []:
            parser = create_object(parser)
            self.parsers[parser.tool_type] = parser
        super().__init__(template, **format_field)

    def format_instruction(self) -> List[dict]:
        """
        格式化指令
        
        返回:
            List[dict]: 格式化后的指令列表
        """
        inst = []
        content = super().format_instruction()
        # 添加系统级指令
        if content.strip():
            msg = dict(role='system', content=content)
            if self.tool_type:
                msg['name'] = self.tool_type
            inst.append(msg)
        # 添加每个解析器的指令
        for name, parser in self.parsers.items():
            content = parser.format_instruction()
            if content.strip():
                inst.append(dict(role='system', content=content, name=name))
        return inst

    def parse_response(self, data: str) -> dict:
        """
        解析响应文本
        
        使用合适的解析器解析响应
        
        参数:
            data (str): 待解析的响应文本
            
        返回:
            dict: 解析结果
        """
        res = dict(tool_type=None, thought=data, action=None, status=ToolStatusCode.NO_TOOL)
        # 尝试使用每个解析器解析响应
        for name, parser in self.parsers.items():
            res = parser.parse_response(data)
            if res['tool_type'] == name:
                break
        return res

    def format_response(self, parsed: dict) -> str:
        """
        格式化响应
        
        使用合适的解析器格式化响应
        
        参数:
            parsed (dict): 解析后的响应字典
            
        返回:
            str: 格式化后的响应文本
        """
        if parsed['action'] is None:
            return parsed['thought']
        assert parsed['tool_type'] in self.parsers
        return self.parsers[parsed['tool_type']].format_response(parsed)
